import torch

from config import device
from utils.sort import sortrows


def tournament_selection(size, pop_size, list_indicator):
    list_indicator = [item.unsqueeze(1) for item in list_indicator]
    combined_indicator = torch.cat(list_indicator, dim=1)
    combined_indicator, loc = torch.unique(combined_indicator, return_inverse=True, dim=0)
    _, rank = sortrows(combined_indicator, return_sorted_indices=True)
    _, rank = torch.sort(rank)
    parents = torch.randint(0, pop_size, (size, pop_size))
    _, best = torch.min(rank[loc[parents]], dim=0)
    parents = parents.reshape(-1)
    return parents[(best + torch.arange(0, pop_size, device=device) * size).to(torch.int)]


if __name__ == '__main__':
    front = torch.tensor([1, 3, 4, 1, 2, 6, 1, 3, 2, 5, 3, 1, 6, 5, 3, 7, 3, 8, 4, 1])
    crowed = torch.tensor(
        [-0.324052404025089, -0.403650295375205, -0.285857956846782, -0.487919223456977, -0.282792116361717,
         -0.390310045783338, -0.234429447962751, -0.300979285771396, -0.185726861943966, -0.635526764607827,
         -0.174727385293618, -0.267502189955617, -0.696734602809150, -0.266253731484487, -0.190088417369015,
         -1.72411372762136, -0.144759258170724, -1.18462614603011, -0.188510737178315, -0.320780694095595])
    print(tournament_selection(2, front.size(0), [front, crowed]))
